Inorder successor in BST

Time: O(H); Space: O(1); medium

Given a binary search tree and a node in it, find the in-order successor of that node in the BST.

If the given node has no in-order successor in the tree, return null.

It’s guaranteed p is one node in the given tree. (You can directly compare the memory address to find p)

Example 1:

1
 \
  2

Input: root = {TreeNode} [1,#,2], p = {TreeNode} [1]

Output: {TreeNode} [2]

Example 2:

  2
 / \
1   3

Input: root = {TreeNode} [2,1,3], p = {TreeNode} [1]

Output: {TreeNode} [2]

Follow up:

  1. O(h), where h is the height of the BST.

[1]:

class TreeNode(object):
    def __init__(self, x):
        self.val = x
        self.left = None
        self.right = None

[2]:

class Solution1(object):
    """
    Time: O(H)
    Space: O(1)
    """
    def inorderSuccessor(self, root, p):
        """
        :type root: TreeNode
        :type p: TreeNode
        :rtype: TreeNode
        """
        # If it has right subtree.
        if p and p.right:
            p = p.right
            while p.left:
                p = p.left
            return p

        # Search from root.
        successor = None
        while root and root != p:
            if root.val > p.val:
                successor = root
                root = root.left
            else:
                root = root.right

        return successor

[6]:

s = Solution1()

root = TreeNode(1)
root.right = TreeNode(2)
p = TreeNode(1)
res = s.inorderSuccessor(root, p)
assert res.val == 2

root = TreeNode(2)
root.left = TreeNode(1)
root.right = TreeNode(3)
p = TreeNode(1)
res = s.inorderSuccessor(root, p)
assert res.val == 2